#include "detector/detector.hpp"
#include "ros/node_handle.h"

namespace detector {

Detector::Detector(const int &light_thresh, const LightParams &light_params,
                   const ArmorParams &armor_params,
                   const std::string &model_path)
    : min_lightness_(light_thresh), l_(light_params), a_(armor_params) {
  classifier_ = std::make_unique<NumberClassifier>(model_path);
}

cv::Mat Detector::preprocessImage(const cv::Mat &src, const int &color) {
  cv::Mat bgr[3];
  cv::split(src, bgr);
  cv::Mat binary;
  cv::threshold(bgr[color], binary, min_lightness_, 255, cv::THRESH_BINARY);
  return binary;
}

std::vector<Light> Detector::findLights(const cv::Mat &rgb_img,
                                        const cv::Mat &binary_img) {
  using std::vector;
  vector<vector<cv::Point>> contours;
  vector<cv::Vec4i> hierarchy;
  cv::findContours(binary_img, contours, hierarchy, cv::RETR_EXTERNAL,
                   cv::CHAIN_APPROX_SIMPLE);

  vector<Light> lights;

  for (const auto &contour : contours) {
    if (contour.size() < 5)
      continue;

    auto r_rect = cv::minAreaRect(contour);
    auto light = Light(r_rect);

    if (isLight(light)) {
      auto rect = light.boundingRect();
      if ( // Avoid assertion failed
          0 <= rect.x && 0 <= rect.width &&
          rect.x + rect.width <= rgb_img.cols && 0 <= rect.y &&
          0 <= rect.height && rect.y + rect.height <= rgb_img.rows) {
        int sum_r = 0, sum_b = 0;
        auto roi = rgb_img(rect);
        // Iterate through the ROI
        for (int i = 0; i < roi.rows; ++i) {
          for (int j = 0; j < roi.cols; ++j) {
            if (cv::pointPolygonTest(
                    contour, cv::Point2f(j + rect.x, i + rect.y), false) >= 0) {
              // if point is inside contour
              sum_r += roi.at<cv::Vec3b>(i, j)[2];
              sum_b += roi.at<cv::Vec3b>(i, j)[0];
            }
          }
        }
        // Sum of red pixels > sum of blue pixels ?
        light.color = sum_r > sum_b ? RED : BLUE;
        lights.emplace_back(light);
      }
    }
  }
  std::sort(lights.begin(), lights.end(),
            [](Light &l1, Light &l2) { return l1.top.x < l2.top.x; });
  return lights;
}

bool Detector::isLight(const Light &light) {
  // The ratio of light (short side / long side)
  float ratio = light.width / light.length;
  bool ratio_ok = l_.min_ratio < ratio && ratio < l_.max_ratio;

  bool angle_ok = light.tilt_angle < l_.max_angle;

  bool is_light = ratio_ok && angle_ok;

  return is_light;
}

std::vector<Armor> Detector::matchLights(const std::vector<Light> &lights,
                                         const int &color) {
  std::vector<Armor> armors;
  // Loop all the pairing of lights
  for (auto light_1 = lights.begin(); light_1 != lights.end(); light_1++) {
    for (auto light_2 = light_1 + 1; light_2 != lights.end(); light_2++) {
      if (light_1->color != color || light_2->color != color)
        continue;

      if (containLight(*light_1, *light_2, lights)) {
        continue;
      }

      auto armor = Armor(*light_1, *light_2);
      if (isArmor(armor)) {
        armors.emplace_back(armor);
      }
    }
  }

  return armors;
}

bool Detector::containLight(const Light &light_1, const Light &light_2,
                            const std::vector<Light> &lights) {
  auto points = std::vector<cv::Point2f>{light_1.top, light_1.bottom,
                                         light_2.top, light_2.bottom};
  auto bounding_rect = cv::boundingRect(points);

  for (const auto &test_light : lights) {
    if (test_light.center == light_1.center ||
        test_light.center == light_2.center) {
      continue;
    }

    if (bounding_rect.contains(test_light.top) ||
        bounding_rect.contains(test_light.bottom) ||
        bounding_rect.contains(test_light.center)) {
      return true;
    }
  }
  return false;
}

bool Detector::isArmor(Armor &armor) {
  Light light_1 = armor.left_light;
  Light light_2 = armor.right_light;
  // Ratio of the length of 2 lights (short side / long side)
  float light_length_ratio = light_1.length < light_2.length
                                 ? light_1.length / light_2.length
                                 : light_2.length / light_1.length;
  bool light_ratio_ok = light_length_ratio > a_.min_light_ratio;

  // Distance between the center of 2 lights (unit : light length)
  float avg_light_length = (light_1.length + light_2.length) / 2;
  float center_distance =
      cv::norm(light_1.center - light_2.center) / avg_light_length;
  bool center_distance_ok = (a_.min_small_center_distance < center_distance &&
                             center_distance < a_.max_small_center_distance) ||
                            (a_.min_large_center_distance <= center_distance &&
                             center_distance < a_.max_large_center_distance);

  // Angle of light center connection
  cv::Point2f diff = light_1.center - light_2.center;
  float angle = std::abs(std::atan(diff.y / diff.x)) / CV_PI * 180;
  bool angle_ok = angle < a_.max_angle;

  bool is_armor = light_ratio_ok && center_distance_ok && angle_ok;
  armor.is_big = center_distance > a_.min_large_center_distance ? true : false;

  return is_armor;
}

std::vector<Armor> Detector::numberClassify(std::vector<Armor> &armors,
                                            cv::Mat &rgb_img) {
  // extract number img
  for (auto &armor : armors) {
    std::vector<cv::Point> points{
        cv::Point{armor.left_light.top}, cv::Point{armor.left_light.bottom},
        cv::Point{armor.right_light.bottom}, cv::Point{armor.right_light.top}};
    cv::RotatedRect rrect = cv::minAreaRect(points);
    if (rrect.size.width > rrect.size.height) {
      rrect.size.height *= 2;
    } else {
      rrect.size.width *= 2;
    }
    cv::Rect roi = rrect.boundingRect();
    // avoid ASSERTION_FAILED
    if (roi.x >= 0 && roi.y >= 0 && roi.x + roi.width <= rgb_img.cols &&
        roi.y + roi.height <= rgb_img.rows && roi.width > 0 && roi.height > 0) {
      cv::Mat tmp = rgb_img(roi);
      cv::resize(tmp, tmp, cv::Size(28, 28));
      cv::cvtColor(tmp, tmp, cv::COLOR_BGR2GRAY);
      armor.number_img = tmp;
    }
  }
  // do classify
  std::vector<Armor> real_armors;
  for (auto &armor : armors) {
    if (armor.number_img.empty()) {
      continue;
    }
    auto [id, confidence] = classifier_->classify(armor.number_img);
    if (id > 0) {
      armor.number = id;
      armor.confidence = confidence;
      real_armors.emplace_back(armor);
    }
  }
  return real_armors;
}

void Detector::set_lightness_thresh(const int &lightness_thresh) {
  min_lightness_ = lightness_thresh;
}

} // namespace detector